#' @TODO 预后效能分析，风险因子联动图和KM ROC曲线合图
#' @title ## 预后效能分析，风险因子联动图和KM ROC曲线合图
#' @param model_coef lasso分析结果，第一列为基因名，第二列为系数数值；
#' @param exp 为表达谱
#' @param dataset 为KM_ROC_res结果中的第三个元素，有sample、riskscore、riskgroup，可以为NULL，则默认使用`KM_ROC_res`中的第三个元素
#' @param KM_ROC_res 为之前函数得到的结果，元素一为KM曲线ggsurplot对象，元素二为ROC曲线是ggplot对象
#' 
#' @examples model_coef数据示例
#' > head(lasso_res[[1]])
#'   symbol        coef
#' 1  SPRY2 -0.01705138
#' 2 CCDC69  0.11278595
#' 3 TUBA1C  0.13265976
#' 4  SFRP1 -0.03003526
#' 5 PSMD14  0.11354158
#' 6   RGS2 -0.05530522
#' 
#' @examples exp数据示例
#' > brca_exp_01[1:4,1:4]
#'         TCGA-A2-A0CY-01A TCGA-B6-A40B-01A TCGA-AO-A0J8-01A TCGA-A8-A08J-01A
#' HEPACAM       0.00000000       0.01093413        0.0178391       0.01105015
#' LEP           0.06739103       0.80698082        0.1018269       0.04282505
#' STARD9        0.23719877       0.40060002        0.2003482       0.30404551
#' ANKRD53       0.13371715       0.58099489        0.1788981       0.20353212
#' 
#' @param KM_ROC_res KM_ROC_curve函数分析结果 
#' @examples KM_ROC_res 数据示例
#' > training
#' [[1]]
#' 
#' [[2]]
#' 
#' [[3]]
#'               sample status         time riskgroup  riskscore
#' 1   TCGA-A2-A0CY-01A      0  55.76666667      High -1.7312376
#' 2   TCGA-B6-A40B-01A      0 105.06666667       Low -3.1421083
#' 3   TCGA-AO-A0J8-01A      0  22.66666667       Low -3.8473514
#' @param output_dir 结果输出目录
#' @param var_name 用于命名文件夹,如果为NULL则使用四个随机字符命名
#' @param surtime_unit 生存时间单位 12或者365对应，月份或者天
#' 
#' @return *list*，多张图，均为gg对象
#' @author *WYK*
#'
Effectiveness_Analysis_v2 <- function(model_coef = NULL, exp = NULL, surtime_unit = c(1, 12, 365),
                                      dataset = NULL, KM_p = NULL, ROC_p = NULL, saveplot = F,
                                      output_dir = NULL, var_name = NULL) {
  if (is.null(var_name)) {
    var_name <- paste0("_", paste0(sample(letters, 4), collapse = ""))
  }

  library(tidyverse)
  library(cowplot)
  library(patchwork)
  library(ggpubr)

  design <- "A
  B
  C
  C
  C
  C"


  p1 <- dataset %>%
    ggplot(mapping = aes(x = reorder(sample, riskscore), y = riskscore)) +
    geom_point(aes(colour = factor(riskgroup, labels = c("High", "Low"))),
      stat = "identity"
    ) +
    theme_test() +
    labs(x = "Patients(increasing score)", y = "Score") +
    theme(
      text = element_text(size = 12),
      plot.margin = unit(c(0.2, 0.2, 0.2, 0.2), "cm"),
      legend.position = "top", axis.ticks.x = element_blank()
    ) +
    guides(colour = guide_legend(title = NULL)) +
    scale_colour_manual(aesthetics = "colour", values = c("High" = "#E41A1C", "Low" = "#367EB8")) +
    theme(legend.position = "left") +
    scale_x_discrete(label = NULL)

  if (surtime_unit == 12) {
    y_chara <- "Survival Time (Month)"
  } else if (surtime_unit == 365) {
    y_chara <- "Survival Time (Day)"
  } else {
    y_chara <- "Survival Time"
  }

  # scales::show_col(ggsci::pal_igv()(12))

  p2 <- ggplot(data = dataset, aes(x = reorder(sample, riskscore), y = time)) +
    geom_point(aes(colour = factor(status, label = c("Alive", "Dead"))), size = 1.4, alpha = .9) +
    theme_test() +
    labs(x = "Patients(increasing score)", y = y_chara) +
    scale_colour_manual(aesthetics = "colour", values = c("#749B58FF", "#802268FF")) +
    # scale_color_manual(labels = c("Alive", "Dead")) +
    guides(color = guide_legend(title = NULL)) +
    theme(
      text = element_text(size = 12),
      plot.margin = unit(c(0.2, 0.2, 0.2, 0.2), "cm"),
      legend.position = "top",
      axis.ticks.x = element_blank()
    ) +
    scale_x_discrete(labels = NULL) +
    theme(legend.position = "left")


  df <- exp[coef_file[[1]], ] %>%
    na.omit() %>%
    as.data.frame() %>%
    DF_z_score(direction = "row")


  df2 <- dataset %>%
    arrange(riskscore) %>%
    mutate(sample = factor(sample, levels = sample))

  df3 <- df %>%
    t() %>%
    as.data.frame() %>%
    rownames_to_column("sample") %>%
    pivot_longer(-sample, values_to = "expr", names_to = "gene") %>%
    mutate(sample = factor(sample, levels = levels(df2[["sample"]])))

  p3_1 <- df2 %>% ggplot(aes(x = sample, y = 1)) +
    geom_col(aes(fill = riskgroup), size = 0, width = 1) +
    scale_fill_manual(values = c("Low" = "#367EB8", "High" = "#E41A1C")) +
    theme_test() +
    theme(
      axis.text.x.bottom = element_blank(),
      axis.ticks.x.bottom = element_blank(),
      axis.line.x.bottom = element_line(color = "black"),
      legend.position = "top", axis.text.y.left = element_blank(),
      axis.ticks.y.left = element_blank()
    ) +
    labs(y = NULL, x = NULL, fill = "Score") +
    scale_x_discrete(labels = NULL, expand = expansion(mult = 0, add = 0)) +
    scale_y_discrete(labels = NULL, expand = expansion(mult = 0, add = 0)) +
    theme(legend.position = "left")

  p3_2 <- df3 %>% ggplot(aes(x = sample, y = gene)) +
    geom_tile(aes(fill = expr)) +
    scale_fill_gradient2(low = muted("navy"), mid = "white", high = muted("firebrick3")) +
    theme_test() +
    theme(
      axis.text.x.bottom = element_blank(), axis.ticks.x.bottom = element_blank(), axis.line.x.bottom = element_line(color = "black"),
      legend.position = "top"
    ) +
    labs(y = NULL, x = NULL, fill = "z-scale") +
    scale_x_discrete(labels = NULL, expand = expansion(mult = 0, add = 0)) +
    scale_y_discrete(expand = expansion(mult = 0, add = 0)) +
    theme(legend.position = "left")

  p3 <- p3_1 / p3_2 +
    plot_layout(heights = c(1, 10), guides = "collect") &
    theme(legend.position = "left")

  p1_3 <- wrap_plots(p1, p2, p3, design = design,tag_level  = 'A')


  km <- ggarrange(KM_p$plot + theme(axis.title.x.bottom = element_blank()),
    KM_p$table,
    ncol = 1, align = "v", heights = c(.7, .31)
  )


  if (!is.null(ROC_p)) {
    kmroc_figure <- plot_grid(km, ROC_p,
      scale = c(1, .965),
      nrow = 2,
      rel_heights = c(1, .97),
      labels = LETTERS[4:5],
      label_size = 20
    )

    library(patchwork)
    lay_out <- "
    111222
    111222
    111222
    111222
    111222
    111222
    "

     p <- plot_grid(p1_3 ,kmroc_figure,rel_widths = c(5,4) ,align = c( "hv"))
    
    p <- p1_3 + kmroc_figure + plot_layout(design = lay_out)

    output_dir <- "/Pub/Users/wangyk/Project_wangyk/Codelib_YK/some_scr"
    dir_name <- sprintf("%soutput/Effectiveness_Analysis/Effectiveness_Analysis_%s", output_dir, var_name)

    if (!dir.exists(dir_name)) {
      dir.create(dir_name, recursive = T)
    }

    if (saveplot) {
      p1_2 <- cowplot::plot_grid(p1, p2, ncol = 1, align = "v") + theme(plot.margin = unit(c(.2, .2, .2, .2), "cm"))
      p_no_heatroc <- cowplot::plot_grid(km, p1_2, ncol = 2, labels = "AUTO", align = "h")

      ggsave2(
        filename = paste0(dir_name, "/Fig_km_scatter.pdf"),
        plot = p_no_heatroc, width = 9.4, height = 5
      )

      ggsave2(
        filename = paste0(dir_name, "/Fig_A-E.pdf"),
        plot = p, width = 10, height = 10
      )

      ggsave2(
        filename = paste0(dir_name, "/FigE.pdf"),
        plot = ROC_p, width = 5, height = 5
      )
    }
  }

  library(cowplot)

  p_scatter <- plot_grid(p1, p2, nrow = 2, align = "v")
  p_no_heat <- plot_grid(km, p_scatter, align = "h", nrow = 1, labels = "AUTO")

  if (saveplot) {
    if (!dir.exists(sprintf("%soutput/Effectiveness_Analysis/Effectiveness_Analysis_%s", output_dir, var_name))) {
      dir.create(sprintf("%soutput/Effectiveness_Analysis/Effectiveness_Analysis_%s", output_dir, var_name), recursive = T)
    } else {
      print(sprintf("Dir '%soutput/Effectiveness_Analysis/Effectiveness_Analysis_%s' is existed.", output_dir, var_name))
    }

    dir_name <- sprintf("%soutput/Effectiveness_Analysis/Effectiveness_Analysis_%s", output_dir, var_name)

    ggsave2(
      filename = paste0(dir_name, "/FigA.pdf"),
      plot = p1, width = 4, height = 2.3
    )

    ggsave2(
      filename = paste0(dir_name, "/FigB.pdf"),
      plot = p2, width = 4, height = 2.3
    )

    ggsave2(
      filename = paste0(dir_name, "/FigC.pdf"),
      plot = p3, width = 4, height = 6
    )

    ggsave2(
      filename = paste0(dir_name, "/FigABC.pdf"),
      plot = p1_3, width = 4, height = 9
    )

    ggsave2(
      filename = paste0(dir_name, "/FigD.pdf"),
      plot = km, width = 5, height = 5
    )

    ggsave2(
      filename = paste0(dir_name, "/Fig_no_heat.pdf"),
      plot = p_no_heat, width = 9, height = 5
    )
  }

  # print(p)

  tmp_list <- list()
  tmp_list[[1]] <- p1
  tmp_list[[2]] <- p2
  tmp_list[[3]] <- p3
  tmp_list[[4]] <- p1_3
  tmp_list[[5]] <- p

  names(tmp_list) <- c("p1", "p2", "p3", "risk_figure", "1_5p")
  return(tmp_list)
}
source("/Pub/Users/wangyk/Project_wangyk/Codelib_YK/some_scr/DF_z_scaled.R")


# 效能分析----------
# load("/Pub/Users/wangyk/project/Poroject/P220215002_F210806004_TNBC_Necroptosis/data/major_res.RData", verbose = T)
# source("/Pub/Users/wangyk/Project_wangyk/Codelib_YK/some_cancers/major_test2.r")
# source("/Pub/Users/wangyk/Project_wangyk/Codelib_YK/some_scr/Effectiveness_Analysis.R")
# suppressMessages(library(tidyverse))
# suppressMessages(library(magrittr))

# load("/Pub/Users/wangyk/project/Poroject/P220215002_F210806004_TNBC_Necroptosis/0.data_prepare/GSE58812.RData")


# vali <- KM_ROC_curve(
#     model_coef = coef_file,
#     exp = GSE58812[['data_exprs']],
#     clinical = GSE58812[['data_clinical']] %>% rename(time = OS.time,status = OS), 
#     surtime_unit = c(365),
#     saveplot = F, output_dir = "./", var_name = NULL, 
#     ROC_time_break = c(1,3, 5), 
#     best_cut = F,
#     manual_cutoff = NULL, do_ROC_CI = F, savetiff = F
# )
# a <- Effectiveness_Analysis_v2(
#     model_coef = coef_file,
#     exp = GSE58812[["data_exprs"]],
#     surtime_unit = c(365),
#     dataset = vali[["ScoreInfor"]], ROC_p = vali[['ROC_p']],KM_p = vali[['KM_p']], saveplot = F , 
#     output_dir = './',
#     var_name = NULL
# )

